#!/bin/bash

# Runs the "345M" parameter model
. path.sh

pip install deps/tiktoken-0.7.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl
cache_dir=~/.cache/huggingface
if [ -d "$cache_dir" ]; then
    rm -rf "$cache_dir"
fi
export CUDA_DEVICE_MAX_CONNECTIONS=1
export NCCL_LAUNCH_TIMEOUT=30
#export NCCL_SOCKET_IFNAME=eth0
#export NCCL_IB_GID_INDEX=3
#export NCCL_IB_HCA=mlx5_2:1
#export NCCL_IB_SL=3
#export NCCL_CHECK_DISABLE=1
export NCCL_P2P_DISABLE=0
#export NCCL_IB_DISABLE=0
#export NCCL_DEBUG=INFO
export NCCL_LL_THRESHOLD=16384
#export NCCL_IB_CUDA_SUPPORT=1

export NVTE_APPLY_QK_LAYER_SCALING=1

DATA_PATH=data/data_v18.16.train.jsonl.${SEQ_LENGTH}.hunyuan.json.ep1.glm4.mg.sample_merge.pt
VALID_DATA_PATH=data/data_v18.16.valid.jsonl.${SEQ_LENGTH}.hunyuan.json.ep1.glm4.mg.no_sample_merge.pt


GPUS_PER_NODE=8
# Change for multinode config
MASTER_ADDR=localhost
MASTER_PORT=12323
NNODES=1
NODE_RANK=0
WORLD_SIZE=$(($GPUS_PER_NODE*$NNODES))

SEQ_LENGTH=8192
SAVE_PREFIX=qwen2.5_mglm_datav19.10
. utils/parse_options.sh


#DATA_PATH=x.jsonl.glm3.sample_merge.pt
#VALID_DATA_PATH=x.jsonl.glm3.sample_merge.pt
#DATA_PATH=$VALID_DATA_PATH

MICRO_BATCH_SIZE=1
GLOBAL_BATCH_SIZE=32
TP=2
PP=2
CP=2
# require to align with weight dimensions
HIDDEN_SIZE=4096
FFN_HIDDEN_SIZE=13696
NUM_LAYERS=40
NUM_HEADS=32
VOCAB_SIZE=151552
LAYERNORM_EPSILON=0.00000015625
######################################

SAVE_PATH=/apdcephfs_qy3/share_976139/users/adrenzhou/exp/exp_meetingGPT/${SAVE_PREFIX}_TP${TP}_PP${PP}_CP${CP}_MBZ${MICRO_BATCH_SIZE}_GBSZ${GLOBAL_BATCH_SIZE}_seq${SEQ_LENGTH}


HF_LLAMA_PATH=/apdcephfs_qy3/share_976139/users/adrenzhou/nlp_workdir/pretrained_models/thudm/glm-4-9b-chat


# TOKENIZER_MODEL=/apdcephfs_qy3/share_976139/users/adrenzhou/nlp_workdir/pretrained_models/Llama-2-7b-chat-hf


#TOKENIZER_MODEL=/apdcephfs_qy3/share_976139/users/adrenzhou/nlp_workdir/pretrained_models/thudm/chatglm3-6b-base/tokenizer.model
TOKENIZER_MODEL=/apdcephfs_qy3/share_976139/users/adrenzhou/nlp_workdir/pretrained_models/thudm/glm-4-9b-chat
TOKENIZER_TYPE=HuggingFaceTokenizer


# CHECKPOINT_PATH=/apdcephfs_qy3/share_976139/users/adrenzhou/nlp_workdir/Megatron-LM/hunyuan-7b-mega-T2P2

#CHECKPOINT_PATH=/apdcephfs_qy3/share_976139/users/adrenzhou/nlp_workdir/Megatron-LM/llama7b-megatron-TP${TP}-PP${PP}
CHECKPOINT_PATH=glm4-megatron-TP${TP}-PP${PP}-TE
#CHECKPOINT_PATH=workspace_sh/glm4-megatron-TP${TP}-PP${PP}-TE
tensorboard_output=$SAVE_PATH/tensorboard
mkdir -p $SAVE_PATH/tensorboard

DISTRIBUTED_ARGS="
    --nproc_per_node $GPUS_PER_NODE \
    --nnodes $NNODES \
    --node_rank $NODE_RANK \
    --master_addr $MASTER_ADDR \
    --master_port $MASTER_PORT
"

GPT_ARGS="
    --tensor-model-parallel-size $TP \
    --pipeline-model-parallel-size $PP \
    --context-parallel-size $CP \
    --sequence-parallel \
    --num-layers $NUM_LAYERS \
    --hidden-size $HIDDEN_SIZE \
    --ffn-hidden-size $FFN_HIDDEN_SIZE \
    --num-attention-heads $NUM_HEADS \
    --seq-length $SEQ_LENGTH \
    --max-position-embeddings $SEQ_LENGTH \
    --micro-batch-size $MICRO_BATCH_SIZE \
    --global-batch-size $GLOBAL_BATCH_SIZE \
    --lr 1e-5 \
    --min-lr 1e-6 \
    --train-iters 500 \
    --lr-decay-iters 400 \
    --lr-decay-style cosine \
    --weight-decay 1e-1 \
    --attention-dropout 0 \
    --hidden-dropout 0 \
    --lr-warmup-iters 0 \
    --clip-grad 1.0 \
    --recompute-activations \
    --recompute-granularity selective \
    --use-distributed-optimizer \
    --optimizer hybridadam \
    --exit-on-missing-checkpoint \
    --no-load-optim \
    --no-load-rng \
    --use-rotary-position-embeddings \
    --normalization RMSNorm \
    --norm-epsilon $LAYERNORM_EPSILON \
    --no-position-embedding \
    --disable-bias-linear \
    --swiglu \
    --no-masked-softmax-fusion \
    --transformer-impl transformer_engine \
    --ckpt-format torch \
    --attention-softmax-in-fp32 \
    --dataloader-type cyclic \
    --tokenizer-type $TOKENIZER_TYPE \
    --tokenizer-model $TOKENIZER_MODEL \
    --untie-embeddings-and-output-weights \
    --apply-query-key-layer-scaling \
    --rotary-interleaved \
    --rotary-percent 0.5 \
    --no-rope-fusion \
    --num-query-groups 2 \
    --make-vocab-size-divisible-by 128 \
    --group-query-attention \
    --rotary-base 5000000 \
    --add-qkv-bias \
    --reset-position-ids \
    --reset-attention-mask \
    --tensorboard-dir $tensorboard_output \
    --bf16
"


    #--use-checkpoint-args \

    #--use-legacy-models \

#    --tokenizer-type Llama2Tokenizer \
 #   --tokenizer-model ${TOKENIZER_MODEL} \

    #--untie-embeddings-and-output-weights \
    # --vocab-size 290943 \
    # --mock-data \
    # --data-path $DATA_PATH \
DATA_ARGS="
    --valid-data-path $VALID_DATA_PATH \
    --train-data-path $DATA_PATH
"

OUTPUT_ARGS="
    --log-interval 1 \
    --save-interval 100 \
    --eval-interval 100 \
    --log-throughput \
    --log-memory-to-tensorboard \
    --log-timers-to-tensorboard \
    --eval-iters 12
"

torchrun $DISTRIBUTED_ARGS pretrain_glm4.py \
    $GPT_ARGS \
    $DATA_ARGS \
    $OUTPUT_ARGS \
    --distributed-backend nccl \
    --save $SAVE_PATH \
    --load $CHECKPOINT_PATH

